import torch
import torchvision.datasets
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision.utils import make_grid

# input = torch.tensor([
#     [1,2,0,3,1],
#     [0,1,2,3,1],
#     [1,2,1,0,0],
#     [5,2,3,1,1],
#     [2,1,0,1,1]
# ],dtype=torch.float32)
# input = torch.reshape(input,(-1,1,5,5))
'''
1.加载CIFAR10数据集，并通过DataLoader批量加载图像数据。
2.定义一个包含MaxPool2d层的神经网络模块，用于对输入图像进行最大池化操作。
3.使用TensorBoard记录输入图像和经过最大池化后的输出图像。
4.遍历数据加载器，对每批数据应用最大池化操作，并将结果写入日志文件。
'''
dataset = torchvision.datasets.CIFAR10("../../data", train=False,
                                       transform=torchvision.transforms.ToTensor(), download=True)
dataLoader = DataLoader(dataset,batch_size=64)

class Module(nn.Module):
    def __init__(self):
        super(Module, self).__init__()
        self.maxpool = torch.nn.MaxPool2d(kernel_size=3, ceil_mode=False)

    def forward(self, x):
        x = self.maxpool(x)
        return x


module = Module()
# input = module(input)
# print(input)
step = 0
writer = SummaryWriter("../../logs_maxpool")
for data in dataLoader:
    imgs, targets = data
    grid_imgs = make_grid(imgs)
    writer.add_image("input", grid_imgs, global_step=step)
    output = module(imgs)
    grid_output = make_grid(output)
    writer.add_image("output", grid_output, global_step=step)
    step += 1
writer.close()
